# -*- coding: utf8 -*-

from sage.structure.parent import Parent
from sage.rings.integer import Integer
from sage.misc.classcall_metaclass import ClasscallMetaclass
from sage.structure.unique_representation import UniqueRepresentation
from sage.sets.disjoint_union_enumerated_sets import DisjointUnionEnumeratedSets
from sage.sets.family import Family
from sage.sets.non_negative_integers import NonNegativeIntegers
from sage.categories.finite_enumerated_sets import FiniteEnumeratedSets
from sage.misc.lazy_attribute import lazy_attribute, lazy_class_attribute
from sage.combinat.binary_tree import LabelledBinaryTree


from sage.combinat.combinat import catalan_number
from sage.rings.arith import binomial

def MBST( packedword ):
    '''
    Create the binary unary tree associate to
    the binary search tree with multiplicity
    '''
    def insert( but, lt, l ):
        BUT = BinaryUnaryTree
        LT = LabelledBinaryTree
        if but.is_empty() :
            return BUT( label = 1 ), LT( [], l )
        if l == lt.label() :
            return BUT( [but[0], but[1]], label = but.label() + 1 ), lt
        elif l < lt.label():
            gbut, glt = insert( but[0], lt[0], l )
            return BUT( [gbut, but[1]], label = but.label() ), LT( [glt, lt[1]], label = lt.label() )
        else:
            dbut, dlt = insert( but[1], lt[1], l )
            return BUT( [but[0], dbut], label = but.label() ), LT( [lt[0], dlt], label = lt.label() )

    but = BinaryUnaryTree()
    lt = LabelledBinaryTree( None )
    reverseword = list( packedword ); reverseword.reverse()
    for l in reverseword :
        but, lt = insert( but, lt, l )
    return but

def but_cmp( bt1, bt2 ):
    return 1 if bt2 in bt1.quasi_quot_greater() else -1

def but_cmp2( bt1, bt2 ):
    return 1 if bt1 in bt2.quasi_quot_smaller() else -1

def but_cmp3( bt1, bt2 ):
    return 1 if bt2 in bt1.greater() else -1

def but_cmp4( bt1, bt2 ):
    return 1 if bt2 in bt1.greater2() else -1

class BinaryUnaryTree( LabelledBinaryTree ):
    __metaclass__ = ClasscallMetaclass

    @staticmethod
    def __classcall_private__( cls, *args, **opts ):
        return cls._auto_parent.element_class( cls._auto_parent, *args, **opts )

    @lazy_class_attribute
    def _auto_parent( cls ):
        return BinaryUnaryTrees_all()

    def __init__( self, parent, children = None, label = None, check = True ):
        if label is not None :
            assert( isinstance( label, ( Integer, int ) ) and label > 0 ), "label must be an non negative integer"
            if children is None or children == [] :
                children = [None, None]
            else:
                LabelledBinaryTree.__init__( self, parent, None )
        LabelledBinaryTree.__init__( self, parent, children, label, check )


    def check( self ):
        assert ( not self or ( len( self ) == 2
            and isinstance( self._label, ( Integer, int ) ) )
        ), "This is not a binary unary tree"

    def size( self ):
        if not self: return 0
        else: return self._label + self[0].size() + self[1].size()

#### Quotient of quasi-permutohedron (with transitive reduction) ####

    def _put_child_at_most_left_in_right( self, k ):
        B = self.parent()._element_constructor_
        if self.is_empty(): return B( label = k )
        else: return B( [self[0]._put_child_at_most_left_in_right( k ), self[1]], label = self.label() )

    def _succ_with_multiplicities( self ):
        B = self.parent()._element_constructor_
        if self.is_empty(): return []
        else:
            res = []
            if self.label() > 1:
                res.append( B( [self[0], self[1]._put_child_at_most_left_in_right( 1 )], label = self.label() - 1 ) )
            return ( res +
                [B( [self[0], d], label = self.label() ) for d in self[1]._succ_with_multiplicities()] +
                [B( [g, self[1]], label = self.label() ) for g in self[0]._succ_with_multiplicities()] )

    def _extract_child_with_one_at_most_right_in_left( self ):
        B = self.parent()._element_constructor_
        def build( bt ):
            if bt.is_empty():
                return False, B()
            elif bt[0].is_empty() and bt[1].is_empty() and bt.label() == 1:
                return True, B()
            else:
                bool, tmp = build( bt[1] )
                return bool, B( [bt[0], tmp], label = bt.label() )
        if self.is_empty(): return []
        bool, bt = build( self[0] )
        if bool:
            return [B( [bt, self[1]], label = self.label() + 1 )]
        else: return []

    def _succ_add_multiplicities( self ):
        B = self.parent()._element_constructor_
        if self.is_empty(): return []
        return ( self._extract_child_with_one_at_most_right_in_left()
            + [B( [self[0], d], label = self.label() ) for d in self[1]._succ_add_multiplicities()]
            + [B( [g, self[1]], label = self.label() ) for g in self[0]._succ_add_multiplicities()] )

    def quasi_quot_succ( self ):
        if self.is_empty(): return []
        return self._succ_add_multiplicities() + self._succ_with_multiplicities()

    def quasi_quot_greater( self ):
        from sage.combinat.tools import transitive_ideal
        return transitive_ideal( lambda x: x.quasi_quot_succ(), self )

    def _put_child_at_most_right_in_left( self, k ):
        B = self.parent()._element_constructor_
        if self.is_empty(): return B( label = k )
        else: return B( [self[0], self[1]._put_child_at_most_right_in_left( k )], label = self.label() )

    def _pred_with_multiplicities( self ):
        B = self.parent()._element_constructor_
        if self.is_empty(): return []
        else:
            res = []
            if self.label() > 1:
                res.append( B( [self[0]._put_child_at_most_right_in_left( 1 ), self[1]], label = self.label() - 1 ) )
            return ( res +
                [B( [self[0], d], label = self.label() ) for d in self[1]._pred_with_multiplicities()] +
                [B( [g, self[1]], label = self.label() ) for g in self[0]._pred_with_multiplicities()] )

    def _extract_child_with_one_at_most_left_in_right( self ):
        B = self.parent()._element_constructor_
        def build( bt ):
            if bt.is_empty():
                return False, B()
            elif bt[0].is_empty() and bt[1].is_empty() and bt.label() == 1:
                return True, B()
            else:
                bool, tmp = build( bt[0] )
                return bool, B( [tmp, bt[1]], label = bt.label() )
        if self.is_empty(): return []
        bool, bt = build( self[1] )
        if bool:
            return [B( [self[0], bt], label = self.label() + 1 )]
        else: return []

    def _pred_add_multiplicities( self ):
        B = self.parent()._element_constructor_
        if self.is_empty(): return []
        return ( self._extract_child_with_one_at_most_left_in_right()
            + [B( [self[0], d], label = self.label() ) for d in self[1]._pred_add_multiplicities()]
            + [B( [g, self[1]], label = self.label() ) for g in self[0]._pred_add_multiplicities()] )

    def quasi_quot_pred( self ):
        if self.is_empty(): return []
        return self._pred_add_multiplicities() + self._pred_with_multiplicities()

    def quasi_quot_smaller( self ):
        from sage.combinat.tools import transitive_ideal
        return transitive_ideal( lambda x: x.quasi_quot_pred(), self )

#####################################################################################################

    def other_succ( self ):
        def label_plus_one( bt ):
            if bt.is_empty(): return 1
            else: return bt.label() + 1
        B = self.parent()._element_constructor_
        if self.is_empty(): return []
        res = []
        if self.label() > 1:
            res.append( B( [
                B( list( self[0] ), label = label_plus_one( self[0] ) ),
                self[1]],
                label = self.label() - 1 ) )
            res.append( B( [
                self[0],
                B( list( self[1] ), label = label_plus_one( self[1] ) )],
                label = self.label() - 1 ) )
        return ( res +
            [B( [g, self[1]], label = self.label() ) for g in self[0].other_succ()] +
            [B( [self[0], d], label = self.label() ) for d in self[1].other_succ()] )

    def greater( self ):
        from sage.combinat.tools import transitive_ideal
        return transitive_ideal( lambda x: x.other_succ(), self )

#####################################################################################################

    def other2_succ( self ):
        B = self.parent()._element_constructor_
        if self.is_empty(): return [] ### ????
        res = []
        if self.label() > 1:
            if self[0].is_empty():
                res += [ B( [B( label = i ) , self[1]], label = self.label() - i ) for i in range( 1, self.label() )]
            if self[1].is_empty():
                res += [ B( [self[0], B( label = i )], label = self.label() - i ) for i in range( 1, self.label() )]
        return ( res +
            [B( [g, self[1]], label = self.label() ) for g in self[0].other2_succ()] +
            [B( [self[0], d], label = self.label() ) for d in self[1].other2_succ()] )

    def greater2( self ):
        from sage.combinat.tools import transitive_ideal
        return transitive_ideal( lambda x: x.other2_succ(), self )

#####################################################################################################

    def f_eq( self ):
        '''
        TESTS::

            sage: LT = BinaryUnaryTree
            sage: lt = LT([LT([],2), None],3); lt
            3[2[., .], .]
            sage: lt.f_eq()
            6
            sage: lt = LT([LT([],2), LT([],4)],3); lt
            3[2[., .], 4[., .]]
            sage: lt.f_eq()
            420
        '''
        from sage.symbolic.ring import SymbolicRing
        from sage.misc.functional import integral
        from sage.functions.other import factorial
        t = SymbolicRing().var( 't' )
        def _tmp( self ):
            if self.is_empty(): return t ** 0
            k = self.label()
            return integral( _tmp( self[0] ) * _tmp( self[1] ) * t ** ( k - 1 ) / factorial( k - 1 ), t )
        return _tmp( self )( t = 1 ) * factorial( self.size() )


class BinaryUnaryTrees( UniqueRepresentation, Parent ):

    @staticmethod
    def __classcall_private__( cls, n = None ):
        if n is None:
            return BinaryUnaryTrees_all()
        else:
            assert( isinstance( n, ( Integer, int ) ) and n >= 0 ), "n must be an non negative integer"
            return BinaryUnaryTrees_size( Integer( n ) )

class BinaryUnaryTrees_all( DisjointUnionEnumeratedSets, BinaryUnaryTrees ):

    def __init__( self ):
        DisjointUnionEnumeratedSets.__init__( 
            self, Family( NonNegativeIntegers(), BinaryUnaryTrees_size ),
            facade = True, keepkey = False )

    def _repr_( self ):
        return "Binary Unary Trees"

    def __contains__( self, x ):
        return isinstance( x, self.element_class )

    def __call__( self, x = None, *args, **keywords ):
        return super( BinaryUnaryTrees, self ).__call__( x, *args, **keywords )

    def _element_constructor_( self, *args, **keywords ):
        return self.element_class( self, *args, **keywords )

    Element = BinaryUnaryTree

class BinaryUnaryTrees_size( BinaryUnaryTrees ):

    def __init__( self, size ):
        super( BinaryUnaryTrees_size, self ).__init__( category = FiniteEnumeratedSets() )
        self._size = size

    def _repr_( self ):
        return "Binary Unary Trees of size %s" % ( self._size )

    def __contains__( self, x ):
        return isinstance( x, self.element_class ) and x.size() == self._size

    def _an_element_( self ):
        return self.first()

    def cardinality( self ):
        n = self._size
        cn = lambda k : catalan_number( k )
        b = lambda k : binomial( n - 1, k - 1 )
        return reduce( lambda x, y : x + y, [b( k ) * cn( k ) for k in range( 1, n + 1 )] )

    def __iter__( self ):
        if self._size == 0:
            yield self._element_constructor_( None, None )
        elif self._size == 1:
            yield self._element_constructor( [None, None], 1 )
        else:
            for i in range( 1, self._size + 1 ):
                # root size : i
                for j in range( self._size - i + 1 ):
                    # left tree size : j
                    for lft in self.__class__( j ):
                        k = self._size - i - j
                        # right tree size : k
                        for rgt in self.__class__( k ):
                            yield self._element_constructor_( [lft, rgt], i )

    @lazy_attribute
    def _parent_for( self ):
        return BinaryUnaryTrees_all()

    @lazy_attribute
    def element_class( self ):
        return self._parent_for.element_class

    def _element_constructor_( self, *args, **keywords ):
        res = self.element_class( self._parent_for, *args, **keywords )
        if res.size() != self._size:
            raise ValueError, "Wrong size"
        return res
